import akshare as ak
import pandas as pd
import numpy as np
import re
from datetime import datetime
import os
import time

# 设置中文显示
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

data_dir = 'data'
start_date = '2022-01-01'
end_date = '2024-12-31'
MAX_RETRIES = 3
RETRY_DELAY = 5


def get_zhonggai_stocks():
    try:
        zhonggai_df = ak.stock_us_spot_em()
        print(f"获取到{len(zhonggai_df)}只美股")
        print("\n数据列名:", zhonggai_df.columns.tolist())
        print("\n数据前5行示例:\n", zhonggai_df.head())

        chinese_stocks = [
            'BABA', 'PDD', 'JD', 'BIDU', 'NIO', 'XPEV', 'LI', 'NTES',
            'TCOM', 'BILI', 'TME', 'VIPS', 'ZTO', 'TAL', 'EDU', 'GDS',
            'IQ', 'KC', 'WB', 'HUYA', 'FUTU', 'DADA', 'YMM', 'DIDI',
            'YUMC', 'BZUN', 'MOMO', 'TIGR', 'RLX', 'HTHT', 'QTT',
            'YY', 'DOYU', 'EH', 'XNET', 'SOHU', 'JRJC', 'SFUN', 'WUBA',
            'RENN', 'ATHM', 'CTRP', 'HX', 'LFC', 'CHL', 'CHU', 'CHA',
            'CNC', 'ZNH', 'CEA', 'GSH', 'ASHR', 'KWEB', 'CQQQ'
        ]

        symbol_column = 'symbol' if 'symbol' in zhonggai_df.columns else '代码'
        print(f"\n使用的股票代码列名: {symbol_column}")

        zhonggai_df['clean_code'] = zhonggai_df[symbol_column].str.extract(r'\.?([A-Z]+)W?$')
        print("\n处理后的部分股票代码:", zhonggai_df['clean_code'].head().tolist())

        zhonggai_df = zhonggai_df[zhonggai_df['clean_code'].isin(chinese_stocks)]
        print(f"筛选出{len(zhonggai_df)}只中概股")
        print("\n筛选出的股票:", zhonggai_df[symbol_column].tolist())

        zhonggai_df['代码'] = zhonggai_df[symbol_column]
        zhonggai_df['名称'] = zhonggai_df['名称']

        for idx, row in zhonggai_df.iterrows():
            try:
                stock_info = ak.stock_us_fundamental_em(symbol=row['代码'])
                if not stock_info.empty:
                    zhonggai_df.loc[idx, 'market_cap'] = stock_info['市值'].values[0]
                    zhonggai_df.loc[idx, 'pe_ratio'] = stock_info['市盈率'].values[0]
            except:
                continue

        return zhonggai_df
    except Exception as e:
        print(f"获取中概股列表失败: {e}")
        return None


def get_stock_data(symbol, start_date, end_date):
    try:
        clean_symbol = symbol.split('.')[-1].rstrip('W')
        print(f"正在获取股票数据: {clean_symbol}")
        stock_df = ak.stock_us_daily(symbol=clean_symbol)
        if not stock_df.empty:
            stock_df['date'] = pd.to_datetime(stock_df['date'])
            stock_df = stock_df[(stock_df['date'] >= start_date) & (stock_df['date'] <= end_date)]
            stock_df.set_index('date', inplace=True)

            stock_df['daily_return'] = stock_df['close'].pct_change()
            stock_df['volatility'] = stock_df['daily_return'].rolling(window=20).std()
            stock_df['turnover_rate'] = stock_df['volume'] / stock_df['volume'].rolling(window=20).mean()

            return stock_df
        return None
    except Exception as e:
        print(f"获取{symbol}历史数据失败: {e}")
        return None


def prepare_daily_data(stock_df, symbol, name):
    if stock_df is None or stock_df.empty:
        return None

    daily_df = stock_df.copy()
    daily_df.reset_index(inplace=True)

    daily_df['symbol'] = symbol
    daily_df['name'] = name

    daily_df['amount'] = daily_df['close'] * daily_df['volume']

    return daily_df[['symbol', 'name', 'date', 'open', 'high', 'low',
                     'close', 'volume', 'amount', 'daily_return', 'volatility']]


def calculate_basic_indicators(df):
    try:
        df['MA20'] = df['close'].rolling(window=20).mean()
        df['MA50'] = df['close'].rolling(window=50).mean()

        df['volatility_20'] = df['daily_return'].rolling(window=20).std()
        df['volatility_50'] = df['daily_return'].rolling(window=50).std()

        df['volume_MA20'] = df['volume'].rolling(window=20).mean()

        print("基本指标计算完成")
        return df
    except Exception as e:
        print(f"计算基本指标失败: {e}")
        return df


def analyze_volatility_factors(stock_df, symbol):
    analysis = {}

    analysis['price_volatility'] = stock_df['close'].std()
    analysis['daily_return_volatility'] = stock_df['daily_return'].std()
    analysis['annualized_volatility'] = stock_df['daily_return'].std() * np.sqrt(252)

    if 'MA20' in stock_df.columns and 'MA50' in stock_df.columns:
        stock_df['ma_cross'] = np.where(stock_df['MA20'] > stock_df['MA50'], 1, -1)
        analysis['ma_cross_frequency'] = (stock_df['ma_cross'].diff() != 0).sum() / len(stock_df)

        analysis['price_ma20_deviation'] = (stock_df['close'] / stock_df['MA20'] - 1).std()
        analysis['price_ma50_deviation'] = (stock_df['close'] / stock_df['MA50'] - 1).std()

    if 'volatility_20' in stock_df.columns:
        analysis['volatility_20_mean'] = stock_df['volatility_20'].mean()
        analysis['volatility_20_std'] = stock_df['volatility_20'].std()

    analysis['volume_volatility'] = stock_df['volume'].std()
    analysis['volume_price_corr'] = stock_df['volume'].corr(stock_df['close'])
    analysis['avg_turnover'] = stock_df['turnover_rate'].mean()

    if 'volume_MA20' in stock_df.columns:
        analysis['volume_trend'] = (stock_df['volume'] / stock_df['volume_MA20']).mean()

    return analysis


def save_all_daily_data(stock_dfs, zhonggai_df):
    stocks_dir = os.path.join(data_dir, 'stocks')
    if not os.path.exists(stocks_dir):
        os.makedirs(stocks_dir)
        print(f"已创建目录: {stocks_dir}")

    for symbol, df in stock_dfs.items():
        name = zhonggai_df[zhonggai_df['代码'] == symbol]['名称'].values[0]
        daily_data = prepare_daily_data(df, symbol, name)
        if daily_data is not None:
            filename = f"{symbol.replace('.', '_')}_daily.csv"
            filepath = os.path.join(stocks_dir, filename)
            daily_data.to_csv(filepath, index=False, encoding='utf-8-sig')
            print(f"已保存 {name} 日线数据到 {filepath}")

    print("所有股票日线数据保存完成")


def filter_date_range(df):
    if not df.empty:
        date_col = None
        if '日期' in df.columns:
            date_col = '日期'
        elif 'date' in df.columns:
            date_col = 'date'
        elif 'release_time' in df.columns:
            date_col = 'release_time'
        if date_col:
            df[date_col] = pd.to_datetime(df[date_col])
            return df[(df[date_col] >= start_date) & (df[date_col] <= end_date)]
    return df


def fetch_data_with_retry(func, *args, **kwargs):
    retries = 0
    while retries < MAX_RETRIES:
        try:
            return func(*args, **kwargs)
        except (ConnectionResetError, ConnectionAbortedError) as e:
            print(f"网络连接错误，重试 {retries + 1}/{MAX_RETRIES}: {e}")
            retries += 1
            time.sleep(RETRY_DELAY)
    print(f"达到最大重试次数，无法获取数据。")
    return None


def collect_relevant_market_factors():
    print("正在收集中概股波动成因相关市场因素数据...")
    factors_dir = os.path.join(data_dir, 'market_factors')
    if not os.path.exists(factors_dir):
        os.makedirs(factors_dir)

    # 美联储利率历史数据
    try:
        fed_rate = fetch_data_with_retry(ak.macro_bank_usa_interest_rate)
        if fed_rate is not None:
            fed_rate = filter_date_range(fed_rate)
            fed_rate.to_csv(os.path.join(factors_dir, 'fed_interest_rate.csv'), encoding='utf-8-sig')
            print("已保存美联储利率历史数据")
    except Exception as e:
        print(f"获取美联储利率数据出错: {e}")

    # 中国出口同比数据
    try:
        cn_trade = fetch_data_with_retry(ak.macro_china_exports_yoy)
        if cn_trade is not None:
            cn_trade = filter_date_range(cn_trade)
            cn_trade.to_csv(os.path.join(factors_dir, 'china_exports_yoy.csv'), encoding='utf-8-sig')
            print("已保存中国出口同比数据")
    except Exception as e:
        print(f"获取中国出口数据出错: {e}")

    # 恒生指数历史数据
    try:
        hang_seng = fetch_data_with_retry(ak.stock_zh_index_daily, symbol="sh000001")
        if hang_seng is not None:
            hang_seng = filter_date_range(hang_seng)
            hang_seng.to_csv(os.path.join(factors_dir, 'hang_seng_index.csv'), encoding='utf-8-sig')
            print("已保存恒生指数历史数据")
    except Exception as e:
        print(f"获取恒生指数数据出错: {e}")


def collect_market_data():
    print("正在获取市场和宏观经济数据...")
    macro_dir = os.path.join(data_dir, 'macro')
    if not os.path.exists(macro_dir):
        os.makedirs(macro_dir)

    # 美国国债收益率数据
    try:
        us_bond = fetch_data_with_retry(ak.bond_zh_us_rate)
        if us_bond is not None:
            us_bond = filter_date_range(us_bond)
            us_bond.to_csv(os.path.join(macro_dir, 'us_bond_yield.csv'))
            print("已保存美国国债收益率数据")
    except Exception as e:
        print(f"获取美国国债收益率数据出错: {e}")

    # LPR历史数据
    try:
        cn_lpr = fetch_data_with_retry(ak.macro_china_lpr)
        if cn_lpr is not None:
            cn_lpr = filter_date_range(cn_lpr)
            cn_lpr.to_csv(os.path.join(macro_dir, 'cn_lpr_history.csv'), encoding='utf-8-sig')
            print("已保存LPR历史数据")
    except Exception as e:
        print(f"获取LPR历史数据出错: {e}")

    # 中国GDP年度数据
    try:
        cn_gdp = fetch_data_with_retry(ak.macro_china_gdp_yearly)
        if cn_gdp is not None:
            cn_gdp = filter_date_range(cn_gdp)
            cn_gdp.to_csv(os.path.join(macro_dir, 'cn_gdp_yearly.csv'), encoding='utf-8-sig')
            print("已保存中国GDP年度数据")
    except Exception as e:
        print(f"获取中国GDP数据出错: {e}")

    # 中国CPI年度数据
    try:
        cn_cpi = fetch_data_with_retry(ak.macro_china_cpi_yearly)
        if cn_cpi is not None:
            cn_cpi = filter_date_range(cn_cpi)
            cn_cpi.to_csv(os.path.join(macro_dir, 'cn_cpi_yearly.csv'), encoding='utf-8-sig')
            print("已保存中国CPI年度数据")
    except Exception as e:
        print(f"获取中国CPI数据出错: {e}")


def main():
    if not os.path.exists(data_dir):
        os.makedirs(data_dir)
        print(f"已创建目录: {data_dir}")

    collect_market_data()
    collect_relevant_market_factors()

    zhonggai_df = get_zhonggai_stocks()
    if zhonggai_df is None or len(zhonggai_df) == 0:
        print("未获取到中概股数据")
        return

    all_results = []
    stock_dfs = {}

    for _, row in zhonggai_df.iterrows():
        symbol = row['代码']
        name = row['名称']
        print(f"\n开始分析 {name} ({symbol})")

        stock_df = get_stock_data(symbol, start_date, end_date)
        if stock_df is None or stock_df.empty or stock_df['close'].isnull().all():
            print(f"跳过 {name} ({symbol}) - 数据无效或为空")
            continue

        stock_df = stock_df.dropna(subset=['close', 'daily_return', 'volatility'])

        stock_dfs[symbol] = stock_df

        stock_df = calculate_basic_indicators(stock_df)

        analysis = analyze_volatility_factors(stock_df, symbol)

        result = {
            'symbol': symbol,
            'name': name,
            'market_cap': row.get('market_cap', 0),
            'pe_ratio': row.get('pe_ratio', 0),
            **analysis
        }
        all_results.append(result)

    results_df = pd.DataFrame(all_results)

    for col in results_df.columns:
        if col not in ['symbol', 'name']:
            results_df[col] = pd.to_numeric(results_df[col], errors='ignore')

    results_df = results_df.sort_values('annualized_volatility', ascending=False)

    results_df.to_csv('zhonggai_volatility_analysis.csv', index=False, encoding='utf-8-sig')

    save_all_daily_data(stock_dfs, zhonggai_df)


if __name__ == "__main__":
    main()